package dkdnserver

import (
	"bufio"
	"dkdns/dkFramework/configs"
	"dkdns/dkFramework/logger"
	"fmt"
	"github.com/miekg/dns"
	"io/ioutil"
	"net"
	"os"
	"os/signal"
	"strings"
	"sync"
	"syscall"
	"time"
)

type CachedResponse struct {
	Msg      *dns.Msg
	ExpireAt time.Time
}

var (
	cache      = make(map[string]CachedResponse)
	cacheMutex sync.RWMutex
)

func init() {
	cache = make(map[string]CachedResponse)
}

// GetCacheMap 返回当前缓存内容（用于调试或监控）
func GetCacheMap() map[string]CachedResponse {
	cacheMutex.RLock()
	defer cacheMutex.RUnlock()
	return cache
}

// isPrivateIP 检查 IP 是否为私有地址
func isPrivateIP(ip net.IP) bool {
	privateCIDRs := []string{
		"10.0.0.0/8",
		"172.16.0.0/12",
		"192.168.0.0/16",
	}

	for _, cidr := range privateCIDRs {
		_, ipNet, err := net.ParseCIDR(cidr)
		if err != nil {
			logger.Println(fmt.Sprintf("Error parsing CIDR: %v", err))
			continue
		}
		if ipNet.Contains(ip) {
			return true
		}
	}
	return false
}

// lookupFromCache 从缓存中查找 DNS 响应
func lookupFromCache(queryKey string) (*dns.Msg, bool) {
	cacheMutex.RLock()
	defer cacheMutex.RUnlock()
	cachedResponse, found := cache[queryKey]
	if !found || time.Now().After(cachedResponse.ExpireAt) {
		return nil, false
	}
	return cachedResponse.Msg.Copy(), true
}

// storeInCache 将 DNS 响应存储到缓存中
func storeInCache(queryKey string, response *dns.Msg, timeout time.Duration) {
	cacheMutex.Lock()
	defer cacheMutex.Unlock()
	cache[queryKey] = CachedResponse{
		Msg:      response.Copy(),
		ExpireAt: time.Now().Add(timeout),
	}
}

// DeleteFromCache 删除指定 key 的缓存项
func DeleteFromCache(queryKey string) {
	cacheMutex.Lock()
	defer cacheMutex.Unlock()
	delete(cache, queryKey)
}

// isIPInRange 检查 IP 是否在指定的 IP 范围内
func isIPInRange(ip string, ipRanges []*net.IPNet) bool {
	parsedIP := net.ParseIP(ip)
	if parsedIP == nil {
		return false
	}

	for _, ipNet := range ipRanges {
		if ipNet.Contains(parsedIP) {
			return true
		}
	}
	return false
}

// loadIPRangesFromFile 从文件中加载 IP 范围
func loadIPRangesFromFile(filepath string) ([]*net.IPNet, error) {
	data, err := ioutil.ReadFile(filepath)
	if err != nil {
		return nil, fmt.Errorf("failed to read IP ranges file: %v", err)
	}

	var ipRanges []*net.IPNet
	lines := strings.Split(string(data), "\n")
	for _, line := range lines {
		line = strings.TrimSpace(line)
		if line == "" {
			continue
		}
		_, ipNet, err := net.ParseCIDR(line)
		if err != nil {
			logger.Println(fmt.Sprintf("Error parsing IP range: %v", err))
			continue
		}
		ipRanges = append(ipRanges, ipNet)
	}
	return ipRanges, nil
}

// performRecursiveDNSQuery 执行递归 DNS 查询
func performRecursiveDNSQuery(r *dns.Msg, recursiveAddr string) (*dns.Msg, time.Duration, error) {
	c := &dns.Client{Timeout: 5 * time.Second} // 设置超时时间
	return c.Exchange(r, recursiveAddr+":53")
}

// getLocalDNSServers 从配置文件中获取本地 DNS 服务器列表
func getLocalDNSServers(appConfigs *configs.Config) ([]string, error) {
	file, err := os.Open(appConfigs.DNS.ServerConf)
	if err != nil {
		return nil, fmt.Errorf("failed to open resolv.conf file: %v", err)
	}
	defer file.Close()

	var servers []string
	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		line := strings.TrimSpace(scanner.Text())
		if strings.HasPrefix(line, "nameserver") {
			fields := strings.Fields(line)
			if len(fields) >= 2 {
				servers = append(servers, fields[1])
			}
		}
	}

	if err := scanner.Err(); err != nil {
		return nil, fmt.Errorf("failed to read resolv.conf file: %v", err)
	}

	if len(servers) == 0 {
		return nil, fmt.Errorf("no DNS servers found in resolv.conf")
	}
	return servers, nil
}

// extractIPAddressesFromResponse 从 DNS 响应中提取 IP 地址
func extractIPAddressesFromResponse(resp *dns.Msg) ([]net.IP, error) {
	var ipAddresses []net.IP
	for _, answer := range resp.Answer {
		if a, ok := answer.(*dns.A); ok {
			ipAddresses = append(ipAddresses, a.A)
		}
	}
	return ipAddresses, nil
}

// handleDNSRequest 处理 DNS 请求
func handleDNSRequest(w dns.ResponseWriter, r *dns.Msg, appConfigs *configs.Config) {
	startTime := time.Now()
	clientIP := w.RemoteAddr().String()

	// 检查是否是有效的查询请求
	if len(r.Question) == 0 {
		logger.Printf("Invalid DNS Query - client: %s, error: no question section", clientIP)
		return
	}

	question := r.Question[0]
	queryName := question.Name
	queryType := dns.TypeToString[question.Qtype]

	// 添加 panic 恢复
	defer func() {
		if err := recover(); err != nil {
			logger.Error("DNS请求处理崩溃", fmt.Errorf("%v", err))
			m := new(dns.Msg)
			m.SetRcode(r, dns.RcodeServerFailure)
			w.WriteMsg(m)
		}
	}()

	ipRanges, err := loadIPRangesFromFile(appConfigs.DNS.Whitelist)
	if err != nil {
		logger.Printf("Error loading IP ranges: %v", err)
		return
	}

	for _, q := range r.Question {
		if q.Qtype != dns.TypeA {
			continue
		}

		name := strings.ToLower(q.Name)
		queryKey := fmt.Sprintf("%s-%d", name, q.Qtype)

		logger.Printf("Processing query for %s from %s", name, clientIP)

		// 检查缓存
		if cachedResponse, found := lookupFromCache(queryKey); found {
			logger.Printf("Cache hit for %s", name)
			resp := new(dns.Msg)
			resp.SetReply(r)
			resp.Answer = cachedResponse.Answer
			resp.Id = r.Id
			w.WriteMsg(resp)
			logElapsedTime(clientIP, startTime, name, "cache hit")
			return
		}

		// 本地 DNS 查询
		localDnsServers, err := getLocalDNSServers(appConfigs)
		if err != nil {
			logger.Error("Error getting local DNS servers", err)
			return
		}

		var localResp *dns.Msg
		for _, localDnsServer := range localDnsServers {
			localResp, _, err = performRecursiveDNSQuery(r, localDnsServer)
			if err == nil {
				break
			}
		}

		if localResp != nil {
			ips, err := extractIPAddressesFromResponse(localResp)
			if err == nil {
				for _, addr := range ips {
					if isPrivateIP(addr) || isIPInRange(addr.String(), ipRanges) {
						storeInCache(queryKey, localResp, appConfigs.DNS.CacheTimeout)
						w.WriteMsg(localResp)
						logElapsedTime(clientIP, startTime, name, "local DNS")
						return
					}
				}
			}
		}

		// 递归 DNS 查询
		recursiveResp, _, err := performRecursiveDNSQuery(r, appConfigs.DNS.RecursiveDNS)
		if err != nil {
			logger.Error("Error performing recursive DNS query", err)
			return
		}
		storeInCache(queryKey, recursiveResp, appConfigs.DNS.CacheTimeout)
		w.WriteMsg(recursiveResp)
		logElapsedTime(clientIP, startTime, name, "recursive DNS")
	}

	// 记录完整的查询信息
	logger.Printf("DNS Query - client: %s, elapsed_ms: %d, query: %s, type: %s, source: %s, status: success",
		clientIP, time.Since(startTime).Milliseconds(), queryName, queryType, "DNS")
}

// logElapsedTime 记录处理耗时
func logElapsedTime(clientIP string, startTime time.Time, name, source string) {
	elapsed := time.Since(startTime).Milliseconds()

	// 添加请求类型和查询状态的判断
	if name == "" {
		name = "<empty_query>"
	}
	if source == "" {
		source = "<unknown_source>"
	}

	logger.Printf("DNS Query - client: %s, elapsed_ms: %d, query: %s, source: %s, type: A",
		clientIP, elapsed, name, source)
}

// cleanupCache 定期清理过期缓存
func cleanupCache() {
	defer func() {
		if err := recover(); err != nil {
			logger.Error("缓存清理崩溃", fmt.Errorf("%v", err))
		}
	}()

	logger.Println("开始清理缓存")
	startTime := time.Now()

	cacheMutex.Lock()
	defer cacheMutex.Unlock()

	count := 0
	for key, cachedResponse := range cache {
		if time.Now().After(cachedResponse.ExpireAt) {
			delete(cache, key)
			count++
		}
	}

	logger.Printf("缓存清理完成，删除了 %d 条记录，耗时：%v", count, time.Since(startTime))
}

// RunDNSServer 启动 DNS 服务器
func RunDNSServer(appConfigs *configs.Config) {
	// 启动一个 goroutine 定期清理缓存
	go func() {
		ticker := time.NewTicker(5 * time.Minute) // 每 5 分钟清理一次
		defer ticker.Stop()
		for {
			select {
			case <-ticker.C:
				cleanupCache()
			}
		}
	}()

	udpHandler := dns.NewServeMux()
	udpHandler.HandleFunc(".", func(w dns.ResponseWriter, r *dns.Msg) {
		handleDNSRequest(w, r, appConfigs)
	})

	server := &dns.Server{
		Addr:    fmt.Sprintf(":%d", appConfigs.DNS.Port),
		Net:     "udp",
		Handler: udpHandler,
	}

	logger.Println(fmt.Sprintf("Starting DNS server at port %d", appConfigs.DNS.Port))
	if err := server.ListenAndServe(); err != nil {
		logger.Println(fmt.Sprintf("Error starting DNS server: %v", err))
		os.Exit(1)
	}

	// 优雅关闭
	signalChan := make(chan os.Signal, 1)
	signal.Notify(signalChan, syscall.SIGINT, syscall.SIGTERM)
	<-signalChan
	logger.Println("Shutting down DNS server")
	server.Shutdown()
}
